import websocket
import json
import ssl
import time
import threading
import redis
import os
from datetime import datetime

pool = redis.ConnectionPool(host='localhost',port=6379, db=7)
r = redis.Redis(connection_pool=pool)

dydx_get_orderbook_interval=os.getenv('DYDX_ORDERBOOK_INTERVAL')
if dydx_get_orderbook_interval:
    dydx_get_orderbook_interval=int(dydx_get_orderbook_interval)
else:
    print("环境变量 'dydx_get_orderbook_interval' 未设置")


dydx_token_str=os.getenv('DYDX_TOKEN')
if dydx_token_str:
    dydx_token_array = dydx_token_str.split(',')
    dydx_token_array = [item.strip() for item in dydx_token_array]
else:
    print("环境变量 'DYDX_TOKEN' 未设置")


def on_error(ws, error):
    print("Error:", error)

def on_close(ws, close_status_code, close_msg):
    print("dydx Connection closed")
    time.sleep(10)
    print("dydx websocket 正在重新连接")
    start_dydx_websocket()

def on_open(ws):
    print("dydx websocket open")

def send_message_periodically(ws,interval):
    while True:
        time.sleep(1)
        if ws.sock and ws.sock.connected:
            for token in dydx_token_array:
                subscription_message = {"type":"subscribe","channel":"v4_orderbook","id":token,"batched":"true"}
                ws.send(json.dumps(subscription_message))
                time.sleep(0.5)
        else:
            print("dydx 断开连接")
            break
        time.sleep(interval-1)

def on_message(ws, message):
    global r
    json_data=json.loads(message)
    if("type" in json_data and json_data["type"]=="subscribed"):
        trading_pair=json_data["id"]
        timestamp = datetime.now().strftime('%Y-%m-%d-%H:%M:%S')
        trade_data = {
                "time": timestamp,
                "trading_pair": trading_pair,
                "buy_order": json_data["contents"]["bids"],
                "sell_order": json_data["contents"]["asks"]
            }
        try:
            r.set(timestamp+"-"+trading_pair+"DYDX-Orderbook", json.dumps(trade_data),ex=86400)
        except (redis.exceptions.ConnectionError, redis.exceptions.TimeoutError) as e:
            print(f"Operation failed: {e}. Attempting to reconnect...")
            time.sleep(5)
            r = redis.Redis(connection_pool=pool)


    if("type" in json_data and json_data["type"]=="channel_batch_data"):
        trading_pair=json_data["id"]
        subscription_message = {"type":"unsubscribe","channel":"v4_orderbook","id":trading_pair,"batched":"true"}
        ws.send(json.dumps(subscription_message))

def start_dydx_websocket():
    websocket_url = "wss://indexer.dydx.trade/v4/ws"
    # 创建WebSocket应用程序对象
    try:
        ws = websocket.WebSocketApp(websocket_url,
                                    on_open=on_open,
                                    on_message=on_message,
                                    on_error=on_error,
                                    on_close=on_close
                                )
        ws_thread = threading.Thread(target=ws.run_forever,kwargs={"sslopt": {"cert_reqs": ssl.CERT_NONE}})
        ws_thread.start()

        time.sleep(3)
        if not ws.sock or not ws.sock.connected:
            raise ConnectionError("dydx WebSocket connection failed")


        send_thread = threading.Thread(target=send_message_periodically, args=(ws, dydx_get_orderbook_interval))
        send_thread.start()

    except Exception as e:
        print(e)